
package net.zschech.gwt.comet.server;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import net.zschech.gwt.comet.server.impl.AtmosphereAsyncServlet;
import java.io.IOException;
import java.io.Serializable;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import javax.servlet.ServletException;
import javax.servlet.http.HttpSession;
import net.zschech.gwt.comet.server.impl.AsyncServlet;
import org.atmosphere.cpr.AtmosphereResource;
import org.atmosphere.cpr.AtmosphereServlet;
import org.atmosphere.cpr.Broadcaster;

/**
 *
 * @author p.havelaar
 */
public class AtmosphereCometServlet extends CometServlet {

    public class GwtAtmosphereResource {

        GwtAtmosphereResource(AtmosphereResource<HttpServletRequest, HttpServletResponse> resource,
                CometServletResponse response) {
            this.cometResponse = response;
            this.atmResource = resource;
        }

        public Broadcaster getBroadcaster() {
            return atmResource.getBroadcaster();
        }

        public HttpSession getSession(boolean create) {
            return atmResource.getRequest().getSession(create);
        }

        public void write(Serializable message) throws IOException {
            cometResponse.write(message);
        }

        public void write(Serializable message, boolean flush) throws IOException {
            cometResponse.write(message, flush);
        }

        public void write(List<? extends Serializable> messages) throws IOException {
            cometResponse.write(messages);
        }

        public void write(List<? extends Serializable> messages, boolean flush) throws IOException {
            cometResponse.write(messages);
        }

        public void broadcast(Serializable message) throws IOException {
            getBroadcaster().broadcast(message, atmResource);
        }

        public void broadcast(List<Serializable> messages) throws IOException {
            getBroadcaster().broadcast(messages, atmResource);
        }

        public AtmosphereResource<HttpServletRequest, HttpServletResponse> getAtmosphereResource() {
            return atmResource;
        }

        public CometServletResponse getCometResponse() {
            return cometResponse;
        }

        public synchronized <T> void setAttribute(String name, T value) {
            attributes.put(name, value);
        }

        public synchronized <T> T getAttribute(String name) {
            return (T) attributes.get(name);
        }

        private CometServletResponse cometResponse;
        private AtmosphereResource<HttpServletRequest, HttpServletResponse> atmResource;
        private Map<String, Object> attributes = new HashMap<String, Object>();
    }

    @Override
    public void init() throws ServletException {
        AtmosphereAsyncServlet async = new AtmosphereAsyncServlet();
        async.init(getServletContext());
        getServletContext().setAttribute(AsyncServlet.SERVLET_CONTEXT_KEY, async);
        super.init();
    }

    @Override
    protected void service(HttpServletRequest request, HttpServletResponse response) throws ServletException, IOException {

        AtmosphereResource<HttpServletRequest, HttpServletResponse> atm
                = (AtmosphereResource<HttpServletRequest, HttpServletResponse>)
                request.getAttribute(AtmosphereServlet.ATMOSPHERE_RESOURCE);
        
        if (atm == null) {
            throw new ServletException("Failed to get atmosphere resource. " +
                    "Did you setup atmosphere.xml correctly?");
        }
        // make sure we are using the correct request & response
        super.service(atm.getRequest(), atm.getResponse());
    }

    @Override
    protected void doComet(CometServletResponse cometResponse) throws ServletException, IOException {
        super.doComet(cometResponse);
        
        doComet(getResource(cometResponse));
    }

    protected void doComet(GwtAtmosphereResource resource) {
    }

    @Override
    public void cometTerminated(CometServletResponse cometResponse, boolean serverInitiated) {
        super.cometTerminated(cometResponse, serverInitiated);
        GwtAtmosphereResource res = getResource(cometResponse);
        cometTerminated(res);
        resources.remove(cometResponse);
    }

    protected void cometTerminated(GwtAtmosphereResource resource) {
    }

    protected GwtAtmosphereResource getResource(CometServletResponse response) {
        GwtAtmosphereResource res = resources.get(response);
        if (res == null) {
            AtmosphereAsyncServlet async = (AtmosphereAsyncServlet) getServletContext()
                .getAttribute(AsyncServlet.SERVLET_CONTEXT_KEY);
            AtmosphereResource atm = async.getAtmosphere(response);
            if (atm != null) {
                res = new GwtAtmosphereResource(atm, response);
                resources.put(response, res);
            }
        }
        return res;
    }

    private static ConcurrentMap<CometServletResponse, GwtAtmosphereResource> resources
            = new ConcurrentHashMap<CometServletResponse, GwtAtmosphereResource>();
}
